import torch
import torch.nn as nn

class RadarCNN(nn.Module):
    def __init__(self):
        super(RadarCNN, self).__init__()
        self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1)  # 输入通道数为3，输出通道数为16
        self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) # 输出通道数为32
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)        # 最大池化
        self.fc1 = nn.Linear(32 * 1 * 1, 64)                    # 全连接层
        self.fc2 = nn.Linear(64, 1)                             # 输出层，回归任务输出单个值
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.conv1(x))  # 卷积1 + 激活
        x = self.pool(self.relu(self.conv2(x)))  # 卷积2 + 激活 + 池化
        x = x.view(x.size(0), -1)  # 展平
        x = self.relu(self.fc1(x))  # 全连接层1 + 激活
        x = self.fc2(x)             # 输出层
        return x

